# Kestrel: An XMPP-based Many-Task Computing Scheduler
# Author: Lance Stout <lancestout@gmail.com>
#
# Copyright 2009-2010 Clemson University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import yaml


class Interface(object):

    def __init__(self, kernel):
        """
        The worker interface listens for XMPP messages that
        contain task requests or other similar administrative
        requests.

        Events listened for by this module:
          message
          task/finished
          task/rejected
          presence/available

        Events raised by this module:
          task/received
          task/cancel
          profile/approved
          message/send
        """
        self.kernel = kernel
        self.events = {'message': self.message,
                       'task/started': self.task_started,
                       'task/finished': self.task_finished,
                       'task/rejected': self.task_rejected,
                       'profile/send': self.send_profile,
                       'presence/available': self.manager_available}
        self.kernel.register(self)
        self.msg_events = {'task': 'task/received',
                           'cancel': 'task/cancel',
                           'profile/approved': 'profile/approved'}

    def message(self, data=None):
        """
        Raise events if a received message has a known format.

        In this particular implementation, messages whose body
        is a dictionary written in YAML or JSON and
        contains a 'type' entry.
        """
        body = data['body']
        body['source'] = data['from']
        try:
            msg_type = body.get('type')
            event_name = self.msg_events.get(msg_type, None)
            if event_name is not None:
                self.kernel.event(event_name, body)
        except:
            pass

    def manager_available(self, agent):
        """
        Send worker profile when manager agent becomes available.
        """
        if agent['bare'] == self.kernel.config('worker/manager/id'):
            self.kernel.event('profile/send',
                              self.kernel.config('worker/profile'))

    def task_started(self, task):
        """Send a task/started message."""
        self.send_task(task, 'task/started')

    def task_finished(self, task):
        """Send a task/finished message."""
        self.send_task(task, 'task/finished', True)

    def task_rejected(self, task):
        """Send a task/rejected message."""
        self.send_task(task, 'task/rejected')

    def send_task(self, task, task_type, output=False):
        """
        Send an XMPP message containing the status of the given task.
        """
        msg = {'type': task_type,
               'task_id': task.task_id,
               'job_id': task.job_id}

        if output:
            msg['stdout'] = task.stdout
            msg['stderr'] = task.stderr

        self.kernel.event('message/send', {'to': task.source,
                                           'body': str(msg)})

    def send_profile(self, profile):
        """
        Send the worker's profile to the manager.
        """
        manager = self.kernel.config('worker/manager/id')
        profile['type'] = 'profile'
        self.kernel.event('message/send', {'to': manager,
                                           'body': str(profile)})
